import { Algorithm } from "../types"

// 数据模型定义
interface MPLPNode {
  id: any
  type: string
  attributes: Record<string, any>
}

interface MPLPEdge {
  source: any
  target: any
  type?: string
  weight?: number
  attributes: Record<string, any>
}

interface MPLPGraph {
  nodes: MPLPNode[]
  edges: MPLPEdge[]
}

// 定义时序图数据结构
interface TemporalGraph {
  nodes: MPLPNode[]
  edges: MPLPEdge[]
  timestamp: string
}

// 修改接口定义
interface THGRLInputParams {
  graph_sequence: any[]
  embedding_dim: number
  time_steps: number
  learning_rate: number
  epochs: number
}

interface NodeEmbedding {
  node_id: string
  embedding: number[]
  node_type?: string
}

interface LinkPrediction {
  source_node: string
  target_node: string
  probability: number
  edge_type?: string
}

interface THGRLOutputParams {
  node_embeddings: NodeEmbedding[]
  link_predictions?: LinkPrediction[]
  algorithm: string
  parameters: Record<string, any>
  training_loss?: number[]
}

interface AlgorithmRequest {
  task_id: string
  task_callback_url?: string
  input_params: {
    thgrl_params: THGRLInputParams
  }
}

// 修改算法实现
export const TemporalHeterogeneousGraphRepresentationLearning: Algorithm = {
  name: "时序异构图表示学习算法(THGRL)",
  description: "基于时间序列的异构图表示学习算法，用于学习图中节点的向量表示并进行链接预测",
  configOptions: [
    {
      id: "embedding_dim",
      label: "嵌入维度",
      type: "number", 
      defaultValue: 64,
      placeholder: "64"
    },
    {
      id: "time_steps",
      label: "时间步数",
      type: "number",
      defaultValue: 5,
      placeholder: "5"
    },
    {
      id: "learning_rate",
      label: "学习率",
      type: "number",
      defaultValue: 0.01,
      placeholder: "0.01"
    },
    {
      id: "epochs",
      label: "训练轮数",
      type: "number", 
      defaultValue: 100,
      placeholder: "100"
    },
    {
      id: "inputData",
      label: "输入数据",
      type: "textarea",
      placeholder: `请输入JSON格式的时序图数据序列，例如：
[
  {
    "nodes": [
      {
        "id": "n1",
        "type": "用户",
        "attributes": {"feature1": 0.8}
      }
    ],
    "edges": [
      {
        "source": "n1",
        "target": "n2",
        "type": "关系",
        "weight": 1.0
      }
    ]
  }
]`
    }
  ],

  execute: async (config) => {
    try {
      const taskId = `thgrl_${Date.now()}_${Math.random().toString(36).substr(2, 9)}`
      
      let graphSequence
      try {
        graphSequence = JSON.parse(config.inputData)
      } catch (e) {
        throw new Error("输入数据必须是有效的JSON格式")
      }

      if (!Array.isArray(graphSequence)) {
        throw new Error("输入数据必须是图序列数组")
      }

      const requestData: AlgorithmRequest = {
        task_id: taskId,
        input_params: {
          thgrl_params: {
            graph_sequence: graphSequence,
            embedding_dim: Number(config.embedding_dim) || 64,
            time_steps: Number(config.time_steps) || 5,
            learning_rate: Number(config.learning_rate) || 0.01,
            epochs: Number(config.epochs) || 100
          }
        }
      }

      // 创建任务
      const createResponse = await fetch("http://172.16.4.151:1122/api/thgrl/create", {
        method: "POST",
        headers: {
          "Content-Type": "application/json",
        },
        body: JSON.stringify(requestData)
      })

      if (!createResponse.ok) {
        const errorData = await createResponse.json()
        throw new Error(errorData.detail || "创建任务失败")
      }

      // 执行任务
      const executeResponse = await fetch(`http://172.16.4.151:1122/api/thgrl/execute/${taskId}`, {
        method: "POST"
      })

      if (!executeResponse.ok) {
        const errorData = await executeResponse.json()
        throw new Error(errorData.detail || "执行任务失败")
      }

      // 轮询获取结果
      let result
      let attempts = 0
      const maxAttempts = 60

      while (attempts < maxAttempts) {
        const getResultResponse = await fetch(`http://172.16.4.151:1122/api/thgrl/result/${taskId}`)
        
        if (!getResultResponse.ok) {
          const errorData = await getResultResponse.json()
          throw new Error(errorData.detail || "获取结果失败")
        }

        result = await getResultResponse.json()
        
        if (result.task_status === "COMPLETED") {
          break
        } else if (result.task_status === "FAILED") {
          throw new Error(result.error_message || "任务执行失败")
        }

        await new Promise(resolve => setTimeout(resolve, 1000))
        attempts++
      }

      if (attempts >= maxAttempts) {
        throw new Error("任务执行超时")
      }

      const outputParams = result.output_params.thgrl_results
      
      // 按照边类型对链接预测结果进行分组
      const groupedPredictions = outputParams.link_predictions?.reduce((acc: Record<string, LinkPrediction[]>, pred: LinkPrediction) => {
        const key = pred.edge_type || 'default'
        if (!acc[key]) {
          acc[key] = []
        }
        acc[key].push(pred)
        return acc
      }, {} as Record<string, LinkPrediction[]>)

      // 生成链接预测结果文本
      const linkPredictionText = groupedPredictions 
        ? Object.entries(groupedPredictions)
          .map(([edgeType, predictions]) => {
            // 确保 predictions 是 LinkPrediction[] 类型
            if (!Array.isArray(predictions)) {
              return `  ${edgeType}类型关系:\n    无有效预测数据`;
            }
            
            const typedPredictions = predictions as LinkPrediction[];
            const topPredictions = typedPredictions
              .sort((a, b) => b.probability - a.probability)
              .slice(0, 5)  // 只显示每种类型概率最高的5个预测
              .map(p => 
                `    ${p.source_node} -> ${p.target_node} (概率: ${(p.probability * 100).toFixed(2)}%)`
              )
              .join('\n');
            
            return `  ${edgeType}类型关系:\n${topPredictions}`;
          })
          .join('\n\n')
        : '无链接预测结果'

      // 生成结果摘要
      const resultText = `时序异构图表示学习完成：
节点嵌入数量: ${outputParams.node_embeddings.length}
嵌入维度: ${outputParams.parameters.embedding_dim}
算法: ${outputParams.algorithm}
训练参数:
  - 时间步数: ${outputParams.parameters.time_steps}
  - 学习率: ${outputParams.parameters.learning_rate}
  - 训练轮数: ${outputParams.parameters.epochs}

链接预测结果:
${linkPredictionText}`

      return {
        result: resultText,
        steps: [],
        executionTime: 0,
        details: {
          node_embeddings: outputParams.node_embeddings,
          link_predictions: outputParams.link_predictions,
          parameters: outputParams.parameters,
          training_loss: outputParams.training_loss
        }
      }
    } catch (e: any) {
      console.error("THGRL算法执行出错:", e)
      throw e
    }
  }
} 